package com.zhangfb95.crosschain.infra.util.tls.server;

import com.zhangfb95.crosschain.infra.enums.SslTypeEnum;
import com.zhangfb95.crosschain.infra.util.StdCertUtil;
import org.springframework.boot.autoconfigure.websocket.servlet.JettyWebSocketServletWebServerCustomizer;
import org.springframework.boot.ssl.SslBundle;
import org.springframework.boot.ssl.SslStoreBundle;
import org.springframework.boot.web.embedded.jetty.JettyServletWebServerFactory;
import org.springframework.boot.web.server.Ssl;
import org.springframework.boot.web.server.WebServer;
import org.springframework.boot.web.servlet.ServletContextInitializer;

import java.io.IOException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;

/**
 * @author zhangfubing
 * @since 2023/12/28
 */
public class JettyWebServerBuilder {

    public static WebServer build(WebServerSslProperties sslProperties, ServletContextInitializer... initializers) {
        JettyServletWebServerFactory webServerFactory = new JettyServletWebServerFactory(sslProperties.getPort());
        setSsl(webServerFactory, sslProperties);
        new JettyWebSocketServletWebServerCustomizer().customize(webServerFactory);
        return webServerFactory.getWebServer(initializers);
    }

    private static void setSsl(JettyServletWebServerFactory webServerFactory, WebServerSslProperties properties) {
        String sslType = properties.getSslType();
        if (SslTypeEnum.NONE.is(sslType)) {
        } else if (SslTypeEnum.in(sslType, SslTypeEnum.RSA, SslTypeEnum.ECDSA)) {
            String defaultBundleName = "default";

            String caCert = properties.getCaCert();
            String sslKey = properties.getSslKey();
            String sslCert = properties.getSslCert();

            try {
                // 构造keyStore
                KeyStore keyStore = KeyStore.getInstance("pkcs12");
                keyStore.load(null);
                // 构造私钥和证书
                PrivateKey privateKey = StdCertUtil.toPrivateKey(sslKey);
                X509Certificate[] certificates = StdCertUtil.toX509Certificates(sslCert);
                // 关联私钥和证书
                keyStore.setKeyEntry("my-key-", privateKey, null, certificates);

                // 构造可信证书源
                KeyStore trustStore = KeyStore.getInstance("pkcs12");
                trustStore.load(null);
                X509Certificate[] caCertificates = StdCertUtil.toX509Certificates(caCert);
                for (int i = 0; i < caCertificates.length; i++) {
                    String alias = "my-trust-" + i;
                    trustStore.setCertificateEntry(alias, caCertificates[i]);
                }

                // 设置ssl绑定
                webServerFactory.setSslBundles(bundleName -> {
                    if (defaultBundleName.equals(bundleName)) {
                        return SslBundle.of(new SslStoreBundle() {
                            @Override
                            public KeyStore getKeyStore() {
                                return keyStore;
                            }

                            @Override
                            public String getKeyStorePassword() {
                                return null;
                            }

                            @Override
                            public KeyStore getTrustStore() {
                                return trustStore;
                            }
                        });
                    } else {
                        throw new RuntimeException("bundleName not find, name: " + bundleName);
                    }
                });

                // 设置ssl
                Ssl ssl = new Ssl();
                ssl.setClientAuth(Ssl.ClientAuth.NEED);
                ssl.setBundle(defaultBundleName);
                webServerFactory.setSsl(ssl);

            } catch (KeyStoreException | IOException | NoSuchAlgorithmException | CertificateException e) {
                throw new RuntimeException(e);
            }
        } else {
            throw new UnsupportedOperationException();
        }
    }
}